import os
import pickle
import argparse

import numpy as np

import DocTer.fuzzer.util as util
from DocTer.fuzzer.fuzzer import Fuzzer
from DocTer.fuzzer.fuzzer_config import FuzzerConfig
from utils.encoder import GreatEncoder


def generate_inputs(constraint_path, dtype_path, output_path, fuzz_optional_p=0.2, guided_mutate=True,
                    mutate_p=0.4, max_iter=10):
    target_config = util.read_yaml(os.path.abspath(constraint_path))
    configs = {
        "fuzz_optional_p": fuzz_optional_p,
        "guided_mutate": guided_mutate,
        "mutate_p": mutate_p,
        "max_iter": max_iter,
        'target_config': target_config,
        'dtype_config': util.read_yaml(os.path.abspath(dtype_path))
    }

    fuzzer_config = FuzzerConfig(configs)
    fuzzer = Fuzzer(fuzzer_config)
    all_name = target_config.get("title")

    full_output_path = os.path.join(output_path, all_name + ".pickle")
    test_cases = []

    with open(full_output_path, "wb+") as f:
        cnt = 0
        while cnt < max_iter:
            new_seed = fuzzer._generate_new_input(True)
            test_cases.append(new_seed)
            cnt += 1
        # write new seeds into pickle
        pickle.dump(test_cases, f)


if __name__ == "__main__":
    current_path = os.path.dirname(os.path.abspath(__file__))
    parser = argparse.ArgumentParser()
    parser.add_argument("--dtype_path", "-d", default=os.path.join(current_path, "DocTer/pytorch_dtypes.yml"))
    parser.add_argument("--constraint_path", "-c", default=os.path.join(current_path, "DocTer/constraints/pytorch"))
    parser.add_argument("--output_path", "-o", default="/home/ubuntu/onnx_samples/onnx_transformer")
    parser.add_argument("--max_iter", type=int, default=1)

    args = parser.parse_args()
    output_path = args.output_path
    constraint_path = args.constraint_path
    # constraint_path_list = os.listdir(constraint_path)
    constraint_path_list = ["torch.nn.functional.conv2d.yaml"]
    for file_name in constraint_path_list:
        api_name = file_name[:-5]
        print(api_name)
        full_constraint_path = os.path.join(constraint_path, file_name)
        full_output_path = os.path.join(output_path, api_name)
        if not os.path.exists(full_output_path):
            os.system("mkdir -p %s" % full_output_path)
        generate_inputs(full_constraint_path, args.dtype_path, full_output_path, max_iter=args.max_iter)
